import os
import csv
import numpy as np


def euler_to_quaternion(roll, pitch, yaw):
    """欧拉角转四元数（Z-Y-X旋转顺序）"""
    cr = np.cos(roll * 0.5)
    sr = np.sin(roll * 0.5)
    cp = np.cos(pitch * 0.5)
    sp = np.sin(pitch * 0.5)
    cy = np.cos(yaw * 0.5)
    sy = np.sin(yaw * 0.5)

    q0 = cr * cp * cy + sr * sp * sy
    q1 = sr * cp * cy - cr * sp * sy
    q2 = cr * sp * cy + sr * cp * sy
    q3 = cr * cp * sy - sr * sp * cy
    return [q0, q1, q2, q3]


def quaternion_rotate_vector(q, v):
    """使用四元数旋转向量（从机体坐标系到惯性坐标系）"""
    q = np.array(q)
    v = np.array(v)

    # 四元数共轭
    q_conj = np.array([q[0], -q[1], -q[2], -q[3]])

    # 构造纯四元数 [0, v]
    v_quat = np.array([0, v[0], v[1], v[2]])

    # 旋转计算：q * v_quat * q_conj
    qv = quaternion_multiply(q, v_quat)
    rotated_quat = quaternion_multiply(qv, q_conj)

    # 提取向量部分（忽略标量部分）
    return rotated_quat[1:]


def quaternion_multiply(q1, q2):
    """四元数乘法"""
    w1, x1, y1, z1 = q1
    w2, x2, y2, z2 = q2

    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2

    return np.array([w, x, y, z])


def compute_total_thrust(row):
    """
    从actuator_thrust和姿态信息计算总推力向量
    参数:
        row: CSV文件中的一行数据
    返回:
        F_thrust_body: 机体坐标系下的总推力向量 [Fx, Fy, Fz]
    """
    try:
        # 获取总推力大小（已考虑姿态控制和电机响应）
        thrust = float(row['actuator_thrust'])

        # 获取姿态信息（优先使用四元数，否则用欧拉角转换）
        try:
            q = [
                float(row['q0']),
                float(row['q1']),
                float(row['q2']),
                float(row['q3'])
            ]
        except KeyError:
            roll = float(row['roll(rad)'])
            pitch = float(row['pitch(rad)'])
            yaw = float(row['yaw(rad)'])
            q = euler_to_quaternion(roll, pitch, yaw)

        # 在机体坐标系中，假设总推力方向沿z轴正方向（垂直向上）
        F_thrust_body = np.array([0.0, 0.0, thrust])

        return F_thrust_body, q

    except KeyError as e:
        raise ValueError(f"数据缺少必要字段：{str(e)}")


def compute_aerodynamic_residual(row, drone_mass=0.5):
    """
    计算气动力残差 f_a（真实值）
    公式：f_a = m·a - F_thrust - F_gravity
    考虑姿态对推力方向的影响
    """
    g = 9.8  # 重力加速度（m/s²）

    # 1. 提取加速度（a）
    try:
        a_x = float(row['acceleration_x(m/s2)'])
        a_y = float(row['acceleration_y(m/s2)'])
        a_z = float(row['acceleration_z(m/s2)'])
    except KeyError:
        raise ValueError("原始数据缺少加速度字段（acceleration_x/y/z(m/s2)）")

    # 2. 计算惯性力（m·a）
    m_dot_a = np.array([
        drone_mass * a_x,
        drone_mass * a_y,
        drone_mass * a_z
    ])

    # 3. 计算总推力向量（考虑姿态）
    F_thrust_body, q = compute_total_thrust(row)

    # 4. 将总推力从机体坐标系转换到惯性坐标系
    F_thrust_inertial = quaternion_rotate_vector(q, F_thrust_body)

    # 5. 计算重力（F_gravity，假设z轴向上为正）
    F_gravity = np.array([0.0, 0.0, -drone_mass * g])

    # 6. 计算气动力残差 f_a
    fa = m_dot_a - F_thrust_inertial - F_gravity

    return fa.tolist()


def convert_single_csv(input_csv, output_csv, drone_mass=0.5):
    """转换单个CSV文件，新增气动力残差计算"""
    try:
        with open(input_csv, 'r') as in_f, open(output_csv, 'w', newline='') as out_f:
            reader = csv.DictReader(in_f)
            fieldnames = ['timestamp(ms)', 'v', 'q', 'pwm', 'fa_true']
            writer = csv.DictWriter(out_f, fieldnames=fieldnames)
            writer.writeheader()

            for row in reader:
                try:
                    # 1. 提取v（速度）
                    v = [
                        float(row['velocity_x(m/s)']),
                        float(row['velocity_y(m/s)']),
                        float(row['velocity_z(m/s)'])
                    ]

                    # 2. 提取姿态
                    try:
                        q = [
                            float(row['q0']),
                            float(row['q1']),
                            float(row['q2']),
                            float(row['q3'])
                        ]
                    except KeyError:
                        roll = float(row['roll(rad)'])
                        pitch = float(row['pitch(rad)'])
                        yaw = float(row['yaw(rad)'])
                        q = euler_to_quaternion(roll, pitch, yaw)

                    # 3. 提取pwm
                    pwm = [
                        float(row['motor1_pwm']),
                        float(row['motor2_pwm']),
                        float(row['motor3_pwm']),
                        float(row['motor4_pwm'])
                    ]

                    # 4. 计算气动力残差（考虑姿态和总推力）
                    fa_true = compute_aerodynamic_residual(row, drone_mass=drone_mass)

                    # 5. 写入数据
                    writer.writerow({
                        'timestamp(ms)': row['timestamp(ms)'],
                        'v': str(v),
                        'q': str(q),
                        'pwm': str(pwm),
                        'fa_true': str(fa_true)
                    })
                except Exception as e:
                    print(f"行 {reader.line_num} 转换失败：{str(e)}")
                    continue

        print(f"✅ 转换完成：{input_csv} -> {output_csv}")
        return True
    except Exception as e:
        print(f"❌ 文件 {input_csv} 转换失败：{str(e)}")
        return False


def batch_convert_csv(input_dir, output_dir, drone_mass=0.5):
    """批量转换目录下的所有CSV文件"""
    os.makedirs(output_dir, exist_ok=True)
    csv_files = [f for f in os.listdir(input_dir) if f.endswith('.csv')]

    if not csv_files:
        print(f"⚠️ 目录 {input_dir} 中没有找到CSV文件")
        return

    print(f"开始批量转换，共 {len(csv_files)} 个文件，无人机质量：{drone_mass}kg")
    success_count = 0

    for csv_file in csv_files:
        input_path = os.path.join(input_dir, csv_file)
        output_file = f"{os.path.splitext(csv_file)[0]}_converted.csv"
        output_path = os.path.join(output_dir, output_file)

        if convert_single_csv(input_path, output_path, drone_mass=drone_mass):
            success_count += 1

    print(f"批量转换完成：成功 {success_count}/{len(csv_files)} 个文件")
    return success_count


if __name__ == '__main__':
    # 配置参数
    INPUT_DIR = "data/raw"  # 原始CSV文件存放目录
    OUTPUT_DIR = "data/processed"  # 转换后的文件存放目录
    DRONE_MASS = 0.029  # 无人机质量（kg）

    batch_convert_csv(INPUT_DIR, OUTPUT_DIR, drone_mass=DRONE_MASS)